import os
import sys
import torch

# 设置项目根目录并切换工作目录
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.chdir(project_root)

# 验证当前工作目录
print(f"Current working directory: {os.getcwd()}")

# 检查 CUDA 是否可用并设置设备
if torch.cuda.is_available():
    try:
        # 尝试初始化 CUDA
        torch.cuda.init()
        device_id = 0  # 使用整数类型的设备 ID
        print(f"使用 GPU: {torch.cuda.get_device_name(0)}")
    except RuntimeError as e:
        print(f"GPU 初始化失败: {e}")
        device_id = 'cpu'
        print("切换到 CPU 模式运行")
else:
    device_id = 'cpu'
    print("警告: 未检测到 GPU，将使用 CPU 模式运行")

import asyncio
import time

from lagent.agents.stream import PLUGIN_CN, get_plugin_prompt
from lagent.distributed import AsyncHTTPAgentClient, AsyncHTTPAgentServer, HTTPAgentClient, HTTPAgentServer
from lagent.llms import INTERNLM2_META
from lagent.schema import AgentMessage
from lagent.utils import create_object

# 创建并设置异步事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# 在创建服务器之前，确保环境变量设置正确
os.environ['JUPYTER_PLATFORM_DIRS'] = '1'

# 在服务器创建之前添加调试信息
print(f"Current working directory: {os.getcwd()}")
print(f"Python path: {sys.path}")

# ----------------------- 测试基础服务器功能 -----------------------
# 创建同步 HTTP 代理服务器
try:
    server = HTTPAgentServer(
        '0' if device_id != 'cpu' else 'cpu',  # 根据设备类型设置 ID
        {
            'type': 'lagent.agents.AsyncAgent',
            'llm': {
                'type': 'lagent.llms.AsyncLMDeployPipeline',
                'path': os.path.abspath('/root/autodl-tmp/models/internlm2_5-1_8b-chat'),
                'meta_template': INTERNLM2_META,
                'tp': 1,
                'pipeline_cfg': {
                    'device': 'cuda:0' if device_id != 'cpu' else 'cpu',  # 明确指定设备
                    'max_batch_size': 1  # 添加批处理大小限制
                }
            }
        },
        port=8090,
    )
    print(server.is_alive)  # 检查服务器状态

    # 测试简单对话
    message = AgentMessage(sender='user', content='hello')
    result = server(message)
    print(result)
finally:
    if 'server' in locals():
        server.shutdown()

# ----------------------- 测试数学解题服务 -----------------------
# 创建异步数学解题服务器
server = AsyncHTTPAgentServer(
    '0' if device_id != 'cpu' else 'cpu',  # 使用动态检测的设备 ID
    {
        'type': 'lagent.agents.AsyncMathCoder',  # 使用异步数学解题代理
        'llm': {
            'type': 'lagent.llms.AsyncLMDeployPipeline',
            'path': '/root/autodl-tmp/models/internlm2_5-1_8b-chat',
            'meta_template': INTERNLM2_META,
            'tp': 1,
            'pipeline_cfg': {
                'device': 'cuda:0' if device_id != 'cpu' else 'cpu',  # 明确指定设备
                'max_batch_size': 1,  # 添加批处理大小限制
                'model_format': 'torch',  # 指定模型格式
                'model_name': 'internlm2-chat-7b',  # 指定模型名称
                'weight_type': 'fp16'  # 使用 fp16 而不是 bf16
            },
            'stop_words': ['<|im_end|>', '<|action_end|>'],
            'max_new_tokens': 1024,
        },
        'interpreter': {
            'type': 'lagent.actions.AsyncIPythonInterpreter',
            'max_kernels': 100
        },
    },
    port=8091,
)

# 测试数学问题求解
message = AgentMessage(
    sender='user',
    content=(
        'Marie is thinking of a multiple of 63, while Jay is thinking of a factor '
        'of 63. They happen to be thinking of the same number. There are two '
        'possibilities for the number that each of them is thinking of, one '
        'positive and one negative. Find the product of these two numbers.'
    )
)
result = server(message)
print(loop.run_until_complete(result))  # 等待异步结果
print(server.state_dict())  # 打印服务器状态

# ----------------------- 测试客户端功能 -----------------------
# 测试异步客户端
client = AsyncHTTPAgentClient(port=8091)
result = client('hello', session_id=1)
print(loop.run_until_complete(result))
print(client.state_dict(1))

# 测试同步客户端
client = HTTPAgentClient(port=8091)
print(client.state_dict(1))
print(client('introduce yourself', session_id=1))
print(client.state_dict(1))
server.shutdown()

# ----------------------- 测试插件服务 -----------------------
# 配置插件和服务器
plugins = [dict(type='lagent.actions.AsyncArxivSearch')]  # 使用异步 Arxiv 搜索插件
server_cfg = dict(
    type='lagent.distributed.AsyncHTTPAgentServer',
    gpu_id='1',
    config={
        'type': 'lagent.agents.AsyncAgentForInternLM',  # 使用异步 InternLM 代理
        'llm': {
            'type': 'lagent.llms.AsyncLMDeployPipeline',
            'path': '/root/autodl-tmp/models/internlm2_5-1_8b-chat',
            'meta_template': INTERNLM2_META,
            'tp': 1,
            'top_k': 1,
            'temperature': 1.0,
            'stop_words': ['<|im_end|>', '<|action_end|>'],
            'max_new_tokens': 1024,
        },
        'plugins': plugins,
        'output_format': {
            'type': 'lagent.prompts.parsers.PluginParser',
            'template': PLUGIN_CN,
            'prompt': get_plugin_prompt(plugins),  # 生成插件提示信息
        }
    },
    port=8091,
)

# 创建服务器并进行并发测试
server = create_object(server_cfg)
tic = time.time()  # 记录开始时间

# 创建 50 个并发请求
coros = [
    server(query, session_id=i)
    for i, query in enumerate(['LLM智能体方向的最新论文有哪些？'] * 50)
]
# 并发执行所有请求
res = loop.run_until_complete(asyncio.gather(*coros))

print('-' * 120)
print(f'time elapsed: {time.time() - tic}')  # 打印总执行时间
server.shutdown()  # 关闭服务器
